#!/usr/bin/env python3
"""
Opendmarc MCP Server

面向 opendmarc 软件的 MCP server，提供 opendmarc 服务管理和命令行操作的工具。

MCP tools 列表：
- opendmarc_service_control: 控制 opendmarc 服务的启动、停止、重启
- opendmarc_service_reload: 重载 opendmarc 服务配置
- opendmarc_service_status: 检查 opendmarc 服务状态

所有工具返回统一的 JSON 结构：
{
  "success": bool,
  "command": str,
  "exit_code": int,
  "stdout": str,
  "stderr": str
}
"""

from mcp.server.fastmcp import FastMCP, Context
import subprocess
import shlex
from typing import Literal

mcp = FastMCP("Opendmarc MCP Server")

@mcp.tool()
def opendmarc_service_control(action: Literal["start", "stop", "restart"]) -> dict:
    """
    控制 opendmarc 服务的启动、停止和重启
    
    参数:
        action: 服务控制动作，可选值为 "start"（启动）、"stop"（停止）、"restart"（重启）
    
    返回:
        JSON 结构: {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    command = f"systemctl {action} opendmarc"
    try:
        result = subprocess.run(
            shlex.split(command),
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def opendmarc_service_reload() -> dict:
    """
    重载 opendmarc 服务配置
    
    返回:
        JSON 结构: {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    command = "systemctl reload opendmarc"
    try:
        result = subprocess.run(
            shlex.split(command),
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def opendmarc_service_status() -> dict:
    """
    检查 opendmarc 服务状态
    
    返回:
        JSON 结构: {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    command = "systemctl status opendmarc"
    try:
        result = subprocess.run(
            shlex.split(command),
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()